package edu.zju.gis.dbfg.model.statistic;

import edu.zju.gis.dbfg.common.exception.ModelFailedException;
import edu.zju.gis.dbfg.model.PyramidConfig;
import edu.zju.gis.dbfg.model.datasource.DataReaderRDD;
import edu.zju.gis.dbfg.model.util.CrsPick;
import edu.zju.gis.dbfg.model.util.GeoItem;
import edu.zju.gis.dbfg.model.util.TileID;
import org.apache.spark.SparkConf;
import org.apache.spark.api.java.JavaPairRDD;
import org.apache.spark.api.java.JavaRDD;
import org.apache.spark.api.java.JavaSparkContext;
import org.apache.spark.api.java.function.*;
import org.geotools.referencing.CRS;
import org.locationtech.jts.geom.Geometry;
import org.opengis.referencing.crs.CoordinateReferenceSystem;
import org.opengis.referencing.operation.MathTransform;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import scala.Tuple2;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.*;

import static edu.zju.gis.dbfg.model.base.Constant.*;

/**
 * 跟据输入的extent图层中的每一个图斑，计算每一个图斑下的各类CC的统计值
 *
 * @author Hu
 */
public class AreaStatistic {

    private final static Logger logger = LoggerFactory.getLogger(AreaStatistic.class);

    public static void main(String[] args) throws Exception {

        // Setup environment
        logger.info("Setup Spark Context");
        SparkConf conf = new SparkConf();
        conf.setMaster("local[4]");
        conf.setAppName("lalala");
        JavaSparkContext jsc = new JavaSparkContext(conf);

        // Setup model arguments
        logger.info("Setup arguments");
        AreaStatisticArgs mArgs = AreaStatisticArgs.initArgs(args);
        if (mArgs == null) {
            String msg = "init args failed, exit";
            logger.error(msg);
            throw new ModelFailedException(AreaStatistic.class, "main()", msg, args);
        }

        CoordinateReferenceSystem crs1 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs1()));
        CoordinateReferenceSystem crs2 = CRS.parseWKT(CrsPick.pickByName(mArgs.getCrs2()));
        CoordinateReferenceSystem targetCRS = DEFAULT_CRS();
        PyramidConfig targetPyramidConfig = DEFAULT_PYRAMID_CONFIG();

        MathTransform transform1 = null;
        MathTransform transform2 = null;
        if (!crs1.equals(targetCRS)) {
            transform1 = CRS.findMathTransform(crs1, targetCRS);
        }
        if (!crs2.equals(targetCRS)) {
            transform2 = CRS.findMathTransform(crs2, targetCRS);
        }

        // Read data from datasource and map it to geoItem
        String input1 = mArgs.getInput1();
        String type1 = mArgs.getType1();
        Map<String, Object> attrs1 = new HashMap<>();
        JavaRDD<GeoItem> gi1 = DataReaderRDD.readToRDD(type1, jsc, input1, null, attrs1, transform1);
        JavaPairRDD<String, GeoItem> gi1Pair = DataReaderRDD.mapToPair(gi1, targetPyramidConfig);

        // Read data from datasource and map it to geoItem
        String input2 = mArgs.getInput2();
        String type2 = mArgs.getType2();
        String pac2 = mArgs.getPacIndex();
        String cc2 = mArgs.getCcIndex();
        String area2 = mArgs.getAreaIndex();
        Map<String, Object> attrs2 = new HashMap<>();
        attrs2.put("CC", cc2);
        attrs2.put("PAC", pac2);
        attrs2.put("AREA", area2);
        JavaPairRDD<String, GeoItem> gi2Pair = null;

        if (type2.equals("ES")) {
            gi1Pair.cache();
            List<TileID> tileIDs = gi1Pair.map(t -> TileID.fromString(t._1)).distinct().collect();
            gi2Pair = DataReaderRDD.readToPairRDD(type2, jsc, input2, attrs2, tileIDs, targetPyramidConfig);
        } else {
            JavaRDD<GeoItem> gi2 = DataReaderRDD.readToRDD(type2, jsc, input2, null, attrs2, transform2);
            gi2Pair = DataReaderRDD.mapToPair(gi2, targetPyramidConfig);
        }

        JavaPairRDD<String, Tuple2<GeoItem, GeoItem>> joined = gi1Pair.join(gi2Pair);
        JavaPairRDD<String, GeoItem> tempRDD = joined.mapToPair(new PairFunction<Tuple2<String, Tuple2<GeoItem, GeoItem>>, String, GeoItem>() {
            @Override
            public Tuple2<String, GeoItem> call(Tuple2<String, Tuple2<GeoItem, GeoItem>> input) throws Exception {
                GeoItem gi2 = input._2._2;
                GeoItem gi1 = input._2._1;
                Geometry g1 = input._2._1.getGeometry();
                Geometry g2 = input._2._2.getGeometry();
                if (!g1.intersects(g2)) {
                    return new Tuple2<>("EMPTY", null);
                }
                Geometry intersection = g1.intersection(g2);
                Double area = intersection.getArea();
                String cc = (String) gi2.getProp().get("CC");
                String pac = String.valueOf(gi2.getProp().get("PAC"));
                String id = pac + "##" + cc;
                GeoItem result = new GeoItem();
                result.setFid(id);
                Map<String, Object> props = new HashMap<>();
                props.put("CC", cc);
                props.put("PAC", pac);
                props.put("AREA", area);
                result.setProp(props);
                result.setGeometry(intersection);
                return new Tuple2<>(id, result);
            }
        }).filter(new Function<Tuple2<String, GeoItem>, Boolean>() {
            @Override
            public Boolean call(Tuple2<String, GeoItem> input) throws Exception {
                return !input._1.equals("EMPTY");
            }
        });

        tempRDD.cache();
        writeToFile(tempRDD, "D:\\ZJU\\test\\FX\\1206\\areaStat.wkt");

        JavaPairRDD<String, GeoItem> resultRDD = tempRDD.reduceByKey(new Function2<GeoItem, GeoItem, GeoItem>() {
            @Override
            public GeoItem call(GeoItem gi1, GeoItem gi2) throws Exception {
                Double area = (Double) gi1.getProp().get("AREA") + (Double) gi2.getProp().get("AREA");
                GeoItem result = new GeoItem();
                result.setFid(gi1.getFid());
                Map<String, Object> props = new HashMap<>();
                props.put("AREA", area);
                result.setProp(props);
                return result;
            }
        });

        Map<String, GeoItem> result = resultRDD.collectAsMap();
        File file = new File(mArgs.getOutput());
        if (file.exists()) {
            file.delete();
        }
        FileWriter writer = new FileWriter(file);
        writer.write(String.format("%s\t%s\t%s\n", pac2, cc2, area2));
        for (String key : result.keySet()) {
            String[] fields = key.split("##");
            String pac = fields[0];
            String cc = fields[1];
            GeoItem gi = result.get(key);
            Double area = (Double) gi.getProp().get("AREA");
            String s = String.format("%s\t%s\t%.9f", pac, cc, area);
            writer.write(s + "\n");
        }
        writer.flush();
        writer.close();
    }

    public static void writeToFile(JavaPairRDD<String, GeoItem> rdd, String filename) throws IOException {
        List<Tuple2<String, GeoItem>> featureResult = rdd.collect();
        File featureFile = new File(filename);
        if (featureFile.exists()) featureFile.delete();
        FileWriter featureWriter = new FileWriter(featureFile);
        for (Tuple2<String, GeoItem> feature : featureResult) {
            GeoItem gi = feature._2;
            String id = gi.getFid();
            Geometry geometry = gi.getGeometry();
            String wkt = geometry.toText();
            featureWriter.write(String.format("%s\t%s\n", id, wkt));
        }
        featureWriter.flush();
        featureWriter.close();
    }

}
